#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAG知识库管理模块
Enhanced RAG Knowledge Base Management Module

功能：
1. 知识库的增删改查
2. 语义搜索
3. 信息验证
4. 知识库更新
"""

import os
import json
import sqlite3
import hashlib
import requests
from typing import Dict, List, Tuple, Optional
from datetime import datetime
import logging
from dataclasses import dataclass
import re

logger = logging.getLogger(__name__)

@dataclass
class KnowledgeEntry:
    """知识库条目"""
    id: str
    title: str
    content: str
    source: str
    source_type: str  # academic, government, news, etc.
    reliability_score: float
    created_at: datetime
    updated_at: datetime
    tags: List[str]
    metadata: Dict

class RAGKnowledgeBase:
    """RAG知识库管理器"""
    
    def __init__(self, db_path: str = "knowledge_base.db"):
        self.db_path = db_path
        self.init_database()
        
    def init_database(self):
        """初始化数据库"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 创建知识库表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS knowledge_entries (
                        id TEXT PRIMARY KEY,
                        title TEXT NOT NULL,
                        content TEXT NOT NULL,
                        source TEXT NOT NULL,
                        source_type TEXT NOT NULL,
                        reliability_score REAL DEFAULT 0.5,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        tags TEXT,
                        metadata TEXT
                    )
                ''')
                
                # 创建全文搜索索引
                cursor.execute('''
                    CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_search 
                    USING fts5(id, title, content, tags)
                ''')
                
                conn.commit()
                logger.info("知识库数据库初始化完成")
                
        except Exception as e:
            logger.error(f"数据库初始化失败: {str(e)}")
    
    def add_entry(self, entry: KnowledgeEntry) -> bool:
        """添加知识库条目"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('''
                    INSERT OR REPLACE INTO knowledge_entries 
                    (id, title, content, source, source_type, reliability_score, 
                     created_at, updated_at, tags, metadata)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    entry.id,
                    entry.title,
                    entry.content,
                    entry.source,
                    entry.source_type,
                    entry.reliability_score,
                    entry.created_at.isoformat(),
                    entry.updated_at.isoformat(),
                    json.dumps(entry.tags),
                    json.dumps(entry.metadata)
                ))
                
                # 更新搜索索引
                cursor.execute('''
                    INSERT OR REPLACE INTO knowledge_search 
                    (id, title, content, tags)
                    VALUES (?, ?, ?, ?)
                ''', (
                    entry.id,
                    entry.title,
                    entry.content,
                    json.dumps(entry.tags)
                ))
                
                conn.commit()
                logger.info(f"成功添加知识库条目: {entry.id}")
                return True
                
        except Exception as e:
            logger.error(f"添加知识库条目失败: {str(e)}")
            return False
    
    def search_entries(self, query: str, limit: int = 10) -> List[KnowledgeEntry]:
        """搜索知识库条目"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 使用全文搜索
                cursor.execute('''
                    SELECT ke.* FROM knowledge_entries ke
                    JOIN knowledge_search ks ON ke.id = ks.id
                    WHERE knowledge_search MATCH ?
                    ORDER BY ke.reliability_score DESC, ke.updated_at DESC
                    LIMIT ?
                ''', (query, limit))
                
                entries = []
                for row in cursor.fetchall():
                    entry = KnowledgeEntry(
                        id=row[0],
                        title=row[1],
                        content=row[2],
                        source=row[3],
                        source_type=row[4],
                        reliability_score=row[5],
                        created_at=datetime.fromisoformat(row[6]),
                        updated_at=datetime.fromisoformat(row[7]),
                        tags=json.loads(row[8]) if row[8] else [],
                        metadata=json.loads(row[9]) if row[9] else {}
                    )
                    entries.append(entry)
                
                return entries
                
        except Exception as e:
            logger.error(f"搜索知识库失败: {str(e)}")
            return []
    
    def verify_claim(self, claim: str) -> Dict:
        """验证声明"""
        try:
            # 搜索相关条目
            entries = self.search_entries(claim, limit=5)
            
            verification_result = {
                "verified": False,
                "contradicts": False,
                "confidence": 0.0,
                "supporting_entries": [],
                "contradicting_entries": [],
                "neutral_entries": []
            }
            
            if not entries:
                return verification_result
            
            # 分析搜索结果
            supporting_count = 0
            contradicting_count = 0
            total_reliability = 0.0
            
            for entry in entries:
                # 简单的关键词匹配验证
                claim_keywords = self._extract_keywords(claim)
                entry_keywords = self._extract_keywords(entry.content)
                
                # 计算相似度
                similarity = self._calculate_similarity(claim_keywords, entry_keywords)
                
                if similarity > 0.7:  # 高相似度
                    if self._check_contradiction(claim, entry.content):
                        verification_result["contradicting_entries"].append(entry)
                        contradicting_count += 1
                    else:
                        verification_result["supporting_entries"].append(entry)
                        supporting_count += 1
                else:
                    verification_result["neutral_entries"].append(entry)
                
                total_reliability += entry.reliability_score
            
            # 计算验证结果
            total_entries = len(entries)
            if total_entries > 0:
                avg_reliability = total_reliability / total_entries
                
                if supporting_count > contradicting_count:
                    verification_result["verified"] = True
                    verification_result["confidence"] = (supporting_count / total_entries) * avg_reliability
                elif contradicting_count > supporting_count:
                    verification_result["contradicts"] = True
                    verification_result["confidence"] = (contradicting_count / total_entries) * avg_reliability
                else:
                    verification_result["confidence"] = avg_reliability * 0.5
            
            return verification_result
            
        except Exception as e:
            logger.error(f"验证声明失败: {str(e)}")
            return {"verified": False, "contradicts": False, "confidence": 0.0}
    
    def _extract_keywords(self, text: str) -> List[str]:
        """提取关键词"""
        # 简单的关键词提取
        keywords = []
        
        # 移除标点符号
        text = re.sub(r'[^\w\s]', '', text)
        
        # 分词（简单按空格分割）
        words = text.lower().split()
        
        # 过滤停用词
        stop_words = {'的', '是', '在', '有', '和', '与', '或', '但', '而', '如果', '因为', '所以', 'the', 'a', 'an', 'and', 'or', 'but', 'if', 'because', 'so'}
        
        for word in words:
            if len(word) > 2 and word not in stop_words:
                keywords.append(word)
        
        return keywords
    
    def _calculate_similarity(self, keywords1: List[str], keywords2: List[str]) -> float:
        """计算关键词相似度"""
        if not keywords1 or not keywords2:
            return 0.0
        
        # 计算Jaccard相似度
        set1 = set(keywords1)
        set2 = set(keywords2)
        
        intersection = len(set1.intersection(set2))
        union = len(set1.union(set2))
        
        return intersection / union if union > 0 else 0.0
    
    def _check_contradiction(self, claim: str, content: str) -> bool:
        """检查是否存在矛盾"""
        # 简单的矛盾检测
        contradiction_indicators = [
            (r'不.*', r'是.*'),
            (r'没有.*', r'有.*'),
            (r'错误.*', r'正确.*'),
            (r'失败.*', r'成功.*'),
            (r'下降.*', r'上升.*'),
            (r'减少.*', r'增加.*')
        ]
        
        claim_lower = claim.lower()
        content_lower = content.lower()
        
        for neg_pattern, pos_pattern in contradiction_indicators:
            if re.search(neg_pattern, claim_lower) and re.search(pos_pattern, content_lower):
                return True
            if re.search(pos_pattern, claim_lower) and re.search(neg_pattern, content_lower):
                return True
        
        return False
    
    def update_entry(self, entry_id: str, updates: Dict) -> bool:
        """更新知识库条目"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 构建更新语句
                set_clause = ", ".join([f"{k} = ?" for k in updates.keys()])
                set_clause += ", updated_at = ?"
                
                values = list(updates.values()) + [datetime.now().isoformat()]
                
                cursor.execute(f'''
                    UPDATE knowledge_entries 
                    SET {set_clause}
                    WHERE id = ?
                ''', values + [entry_id])
                
                if cursor.rowcount > 0:
                    conn.commit()
                    logger.info(f"成功更新知识库条目: {entry_id}")
                    return True
                else:
                    logger.warning(f"未找到要更新的条目: {entry_id}")
                    return False
                    
        except Exception as e:
            logger.error(f"更新知识库条目失败: {str(e)}")
            return False
    
    def delete_entry(self, entry_id: str) -> bool:
        """删除知识库条目"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('DELETE FROM knowledge_entries WHERE id = ?', (entry_id,))
                cursor.execute('DELETE FROM knowledge_search WHERE id = ?', (entry_id,))
                
                conn.commit()
                
                if cursor.rowcount > 0:
                    logger.info(f"成功删除知识库条目: {entry_id}")
                    return True
                else:
                    logger.warning(f"未找到要删除的条目: {entry_id}")
                    return False
                    
        except Exception as e:
            logger.error(f"删除知识库条目失败: {str(e)}")
            return False
    
    def get_statistics(self) -> Dict:
        """获取知识库统计信息"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 总条目数
                cursor.execute('SELECT COUNT(*) FROM knowledge_entries')
                total_entries = cursor.fetchone()[0]
                
                # 按来源类型统计
                cursor.execute('''
                    SELECT source_type, COUNT(*) 
                    FROM knowledge_entries 
                    GROUP BY source_type
                ''')
                source_type_stats = dict(cursor.fetchall())
                
                # 平均可靠性评分
                cursor.execute('SELECT AVG(reliability_score) FROM knowledge_entries')
                avg_reliability = cursor.fetchone()[0] or 0.0
                
                # 最近更新时间
                cursor.execute('SELECT MAX(updated_at) FROM knowledge_entries')
                last_updated = cursor.fetchone()[0]
                
                return {
                    "total_entries": total_entries,
                    "source_type_distribution": source_type_stats,
                    "average_reliability": avg_reliability,
                    "last_updated": last_updated
                }
                
        except Exception as e:
            logger.error(f"获取统计信息失败: {str(e)}")
            return {}

class KnowledgeBaseManager:
    """知识库管理器"""
    
    def __init__(self, db_path: str = "knowledge_base.db"):
        self.kb = RAGKnowledgeBase(db_path)
        
    def add_academic_paper(self, title: str, content: str, source: str, 
                          reliability_score: float = 0.9) -> bool:
        """添加学术论文"""
        entry = KnowledgeEntry(
            id=self._generate_id(title, source),
            title=title,
            content=content,
            source=source,
            source_type="academic",
            reliability_score=reliability_score,
            created_at=datetime.now(),
            updated_at=datetime.now(),
            tags=["academic", "research"],
            metadata={"type": "paper", "reliability": reliability_score}
        )
        
        return self.kb.add_entry(entry)
    
    def add_government_document(self, title: str, content: str, source: str,
                               reliability_score: float = 0.95) -> bool:
        """添加政府文件"""
        entry = KnowledgeEntry(
            id=self._generate_id(title, source),
            title=title,
            content=content,
            source=source,
            source_type="government",
            reliability_score=reliability_score,
            created_at=datetime.now(),
            updated_at=datetime.now(),
            tags=["government", "official"],
            metadata={"type": "document", "reliability": reliability_score}
        )
        
        return self.kb.add_entry(entry)
    
    def add_news_article(self, title: str, content: str, source: str,
                         reliability_score: float = 0.7) -> bool:
        """添加新闻文章"""
        entry = KnowledgeEntry(
            id=self._generate_id(title, source),
            title=title,
            content=content,
            source=source,
            source_type="news",
            reliability_score=reliability_score,
            created_at=datetime.now(),
            updated_at=datetime.now(),
            tags=["news", "media"],
            metadata={"type": "article", "reliability": reliability_score}
        )
        
        return self.kb.add_entry(entry)
    
    def _generate_id(self, title: str, source: str) -> str:
        """生成条目ID"""
        content = f"{title}_{source}_{datetime.now().isoformat()}"
        return hashlib.md5(content.encode()).hexdigest()
    
    def verify_information(self, content: str) -> Dict:
        """验证信息"""
        return self.kb.verify_claim(content)
    
    def get_statistics(self) -> Dict:
        """获取统计信息"""
        return self.kb.get_statistics()

def main():
    """演示知识库使用"""
    manager = KnowledgeBaseManager()
    
    # 添加示例数据
    print("=== 添加示例知识库条目 ===")
    
    # 添加学术论文
    academic_content = """
    人工智能在教育领域的应用研究显示，AI辅助教学可以显著提高学生的学习效率。
    根据最新的研究数据，使用AI技术的学校在标准化测试中的表现比传统学校高出15-30%。
    专家认为，AI技术将在未来十年内彻底改变教育模式。
    """
    
    success = manager.add_academic_paper(
        title="AI在教育领域的应用研究",
        content=academic_content,
        source="https://arxiv.org/abs/2023.12345",
        reliability_score=0.9
    )
    print(f"添加学术论文: {'成功' if success else '失败'}")
    
    # 添加政府文件
    gov_content = """
    教育部发布关于推进教育信息化的指导意见，明确提出要大力发展人工智能教育应用。
    文件指出，到2025年，全国中小学AI教育普及率要达到80%以上。
    政府将投入专项资金支持AI教育基础设施建设。
    """
    
    success = manager.add_government_document(
        title="教育部关于推进教育信息化的指导意见",
        content=gov_content,
        source="http://www.moe.gov.cn/policy/2023/ai_education.html",
        reliability_score=0.95
    )
    print(f"添加政府文件: {'成功' if success else '失败'}")
    
    # 验证信息
    print("\n=== 验证信息 ===")
    test_claim = "AI技术可以提高学生学习效率30%"
    verification_result = manager.verify_information(test_claim)
    
    print(f"验证声明: {test_claim}")
    print(f"验证结果: {verification_result}")
    
    # 显示统计信息
    print("\n=== 知识库统计 ===")
    stats = manager.get_statistics()
    print(f"统计信息: {stats}")

if __name__ == "__main__":
    main() 